一款清华自研的深度学习框架，一键转换PyTorch

点击蓝字默默关注

整理 | 叶子

一、软件介绍

计图（Jittor）是一个完全基于动态编译（Just-in-time），内部使用创新的元算子和统一计算图的深度学习框架。元算子和 Numpy 一样易于使用，而统一计算图则是融合了静态计算图和动态计算图的诸多优点，在易于使用的同时，提供高性能的优化。基于元算子开发的深度学习模型，可以被计图实时的自动优化并且运行在指定的硬件上，如 CPU、GPU。

二、项目地址

https://gitee.com/mirrors/jittor

三、特性

• 算子动态编译：Jittor 内置元算子编译器，可以将用户通过元算子编写的 Python 代码，动态编译成高性能的C++ 代码。

• 自动优化：jittor 内置优化编译遍（complier pass），同时和 LLVM 兼容，这些编译遍会根据硬件设备自动优化动态编译的代码，常见的优化编译遍有循环重排、循环分裂、循环融合、数据打包、向量化、GPU并行。这些编译遍对 C++ 代码进一步优化，生成对计算设备友好的底层算子。

• 统一内存管理：Jittor 使用了统一内存管理，统一 GPU 与 CPU 之间的内存，当深度学习模型将 GPU 内存耗尽时，将使用 CPU 内存来弥补。

• 高效同步异步接口：Jittor 同时提供两种接口，同步和异步接口之间切换不会产生性能损失。

• 模型迁移：Jittor 采用和 PyTorch 相似的模块化接口，同时提供辅助转换脚本，可将 PyTorch 的代码转换成 Jittor 的模型，并可以相互加载和调用。

• 元算子融合：可将神经网络的基本算子（元算子）融合成复杂算子，如卷积层、归一化层，进一步构成神经网络和深度学习应用。

• 高阶导数以及反向传播闭包：元算子是反向传播闭包（元算子反向传播也是元算子），同时支持计算任意高阶导数。在深度学习算子开发过程中，免去反向传播算子重复开发工作，同时可以使用统一的优化策略。

四、安装

Jittor前端语言为Python，使用了模块化的设计，类似于PyTorch，Keras；后端则使用高性能语言编写，如CUDA，C++。当前Jittor支持三种即时编译器：

• CPU 编译器 （需要下列至少一个）

• g++ （>=5.4.0）

• clang （>=8.0）推荐

• GPU 编译器（可选）

• nvcc（>=10.0）

五、示例代码

Jittor 前端语言为 Python。前端使用模块化的设计，类似于 PyTorch，Keras，后端则使用高性能语言编写，如 CUDA，C++。Jittor 官网目前已提供了示例代码，我们可以来分析一下。

下面的代码演示了如何一步一步使用 Python 代码，从头对一个双层神经网络建模。

import jittor as jtfrom jittor import Modulefrom jittor import nn
class Model(Module):   def __init__(self):       self.layer1 = nn.Linear(1, 10)       self.relu = nn.Relu()       self.layer2 = nn.Linear(10, 1)   def execute (self,x) :       x = self.layer1(x)       x = self.relu(x)       x = self.layer2(x)       return x
def get_data(n): # generate random data for training test. for i in range(n):       x = np.random.rand(batch_size, 1)       y = x*x yield jt.float32(x), jt.float32(y)
model = Model()learning_rate = 0.1optim = nn.SGD(model.parameters(), learning_rate)
for i,(x,y) in enumerate(get_data(n)):   pred_y = model(x)   loss = ((pred_y - y)**2)   loss_mean = loss.mean()   optim.step (loss_mean)   print(f"step {i}, loss = {loss_mean.data.sum()}")

参考资料：

https://my.oschina.net/charmve/blog/4472193

往期推荐

扫码关注最新动态

公众号ID:fosslab

我就知道你“在看”